/*
 * batch_main.cpp
 *
 *  Created on: Feb 28, 2011
 *      Author: tqlong
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
#include <armadillo>
#include <ctime>

class ArmaVec : public arma::vec
{
public:
  ArmaVec(int d = 1) : arma::vec(d) {}
  ArmaVec& operator=(const ArmaVec&v) { arma::vec::operator=(v); return *this; }
  double get(int i) const { return at(i); }
  double& ref(int i) { return at(i); }
  int dim() const { return arma::vec::n_elem; }
};

#include "containers/sparse_vec.h"
#include "data/pointset.h"
#include "covertree.h"
#include "kdtree.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);

template <typename T>
void treeBuild()
{
  typedef T                                 node_type;
  typedef typename node_type::pset_type     pset_type;
  typedef typename node_type::iset_type     iset_type;
  typedef typename node_type::vec_type      vec_type;
//  typedef typename node_type::coverset_type coverset_type;

  pset_type data;
  data.load(vm["data"].as<std::string>());
  std::cout << "Done loading data\n"; std::cout.flush();
  pset_type query_data;
  query_data.load(vm["query"].as<std::string>());
  std::cout << "Done loading query\n"; std::cout.flush();

  assert(query_data.dim() == data.dim());

  iset_type idx;
  node_type root(data, idx);
  root.prepareRoot();
  root.splitRecursive();

//  root.print();

  std::vector<int> minIdx, queryIdx;
  node_type query(query_data, queryIdx);
  query.prepareRoot();
  query.splitRecursive();

  std::cout << "Trees are built\n";
  node_type::allNearest(query, root, minIdx);
  std::cout << "all NN done\n";
//  for (unsigned int i = 0; i < minIdx.size(); i++) {
//    int qIdx = queryIdx[i], rIdx = idx[minIdx[i]];
//    std::cout << i << " <--> " << minIdx[i] << " : "
//              << qIdx << " <--> " << rIdx << " d = "
//              << arma::norm(query_data.x(qIdx) - data.x(rIdx),2)
//              << "\n";
//  }
//
  for (unsigned int i = 0; i < minIdx.size(); i++) {
    int qIdx = queryIdx[i], rIdx = idx[minIdx[i]];
    double dist = arma::norm(query_data.x(qIdx) - data.x(rIdx),2);
    bool ok = true;
    for (int j = 0; j < data.n(); j++)
      if (dist > arma::norm(query_data.x(qIdx)-data.x(j),2)) ok = false;

    if (!ok) {
      std::cout << i << " " << qIdx << " " << rIdx << " NOT OK\n";
      return;
    }
  }
  std::cout << "ALL OK !!!\n";
}

template <typename vec>
void generateData()
{
  double noise = vm["noise"].as<double>();
  int n = vm["n"].as<int>();
  int d = vm["d"].as<int>();
  PointSet<vec> data;
  data.generate(n, d, noise);
  data.save(vm["data"].as<std::string>());
}

int main(int argc, char** argv) {
//  test_sparse(); return 0;
//  test_ublas(); return 0;
  process_options(argc, argv);
  if (vm.count("help") || !vm.count("data")) {
    std::cout << desc << "\n";
    return 0;
  }
  if (vm.count("gen")) {
    if (vm.count("seed"))
      srand(time(NULL));
    if (vm.count("sparse"))
      generateData <arma::sparse_vec> ();
    else
      generateData <ArmaVec> ();
    return 0;
  }

  if (!vm.count("query")) {
    std::cout << desc << "\n";
    return 0;
  }

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());
  if (vm.count("sparse")) {
    if (vm["method"].as<std::string>() == "covertree")
      treeBuild<CoverTreeNode<arma::sparse_vec> > ();
    else
      treeBuild<KDTreeNode<arma::sparse_vec, BoundingBox<arma::sparse_vec> > > ();
  }
  else {
    if (vm["method"].as<std::string>() == "covertree")
      treeBuild<CoverTreeNode<ArmaVec> > ();
    else
      treeBuild<KDTreeNode<ArmaVec, BoundingBox<ArmaVec> > > ();
  }
  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  std::cerr << "Duration = " << duration << "\n";
  return 0;
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("data", po::value<std::string>(), "file contains training samples")
    ("query", po::value<std::string>(), "file contains training queries")
    ("sparse", "use sparse vector representation")
    ("gen", "generate random data")
    ("seed", "using different random seed")
    ("noise", po::value<double>()->default_value(0.2), "learning rate")
    ("n", po::value<int>()->default_value(10), "number of points")
    ("d", po::value<int>()->default_value(2), "input dimension")
    ("iter", po::value<int>()->default_value(10), "maximum number of iterations")
    ("alpha", po::value<double>()->default_value(1.0), "learning rate")
    ("C", po::value<double>()->default_value(100.0), "regularized parameter")
    ("method", po::value<std::string>()->default_value("covertree"), "tree algorithm: covertree, kdtree")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

//  root.print();

//  vec_type v(data.dim());
//  v[0] = 5;
//  int nn = nearest<vec_type, node_type >(v, root);
//  std::cout << "nn = " << nn << " idx[nn] = " << idx[nn] << "\nvec =\n" << data.x(idx[nn]) << "\n";

//  for (int i = 0; i < query.n(); i++) {
//    const vec_type& v = query.x(i);
//    int nn = nearest<vec_type, node_type >(v, root);
//    std::cout << i << " " << idx[nn] << "\n";
//  }
